package com.gitee.yanfanvip.single;

import java.io.File;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.function.BiConsumer;
import org.rocksdb.ColumnFamilyDescriptor;
import org.rocksdb.ColumnFamilyHandle;
import org.rocksdb.ColumnFamilyOptions;
import org.rocksdb.DBOptions;
import org.rocksdb.FlushOptions;
import org.rocksdb.InfoLogLevel;
import org.rocksdb.Logger;
import org.rocksdb.Options;
import org.rocksdb.ReadOptions;
import org.rocksdb.RocksDB;
import org.rocksdb.RocksDBException;
import org.rocksdb.RocksIterator;
import org.rocksdb.Snapshot;
import org.rocksdb.WriteBatch;
import org.rocksdb.WriteOptions;
import com.gitee.yanfanvip.interfaces.Database;

public class SingleDatabase implements Database{
	
	static {
		RocksDB.loadLibrary();
	}
	
	static String DATABASE_PATH = "DATABASE";
	final Options OPTIONS = new Options();
	
	final DBOptions DB_OPTIONS = new DBOptions()
			.setCreateIfMissing(true)
			.setCreateMissingColumnFamilies(true)
			.setMaxOpenFiles(500)
			.setMaxTotalWalSize(10000)
			.setLogger(new Logger(OPTIONS) {
				@Override
				protected void log(InfoLogLevel infoLogLevel, String logMsg) {}
			});
	
	final ColumnFamilyOptions COLUMN_FAMILY_OPTIONS = new ColumnFamilyOptions();
	
	final ReadOptions READ_OPTIONS = new ReadOptions().setFillCache(false);
	final WriteOptions WRITE_OPTIONS = new WriteOptions().setDisableWAL(false).setSync(false);
	final HashMap<String, ColumnFamilyHandle> columnFamilys = new HashMap<String, ColumnFamilyHandle>();
	RocksDB db;
	
	static SingleDatabase database = null;
	
	public static synchronized SingleDatabase get(String path) throws RocksDBException {
		if(database != null) {
			return database;
		}else {
			database = new SingleDatabase(DATABASE_PATH + "/" +  path);
			Runtime.getRuntime().addShutdownHook(new Thread() {
				@Override
				public void run() { database.close(); }
			});
			return database;
		}
	}
	
	private SingleDatabase(String path) throws RocksDBException {
		if(!new File(path).exists()){ new File(path).mkdirs(); }
		List<byte[]> cfs = columnFamilys(path);
		List<ColumnFamilyDescriptor> cfds = new ArrayList<ColumnFamilyDescriptor>();
		cfs.forEach(cf->{
			cfds.add(new ColumnFamilyDescriptor(cf, COLUMN_FAMILY_OPTIONS));
		});
		if(cfds.size() == 0) { cfds.add(new ColumnFamilyDescriptor("default".getBytes())); }
		List<ColumnFamilyHandle> cfhs = new ArrayList<ColumnFamilyHandle>();
		db = RocksDB.open(DB_OPTIONS, path, cfds, cfhs);
		for (int j = 0; j < cfs.size(); j++) {
			byte[] bs = cfs.get(j);
			columnFamilys.put(new String(bs), cfhs.get(j));
		}
	}
	
	private static List<byte[]> columnFamilys(String path) throws RocksDBException {
		try(Options options = new Options()){
			return RocksDB.listColumnFamilies(options, path);
		}
	}
	
	ColumnFamilyHandle getColumnFamily(String columnFamilyName) throws RocksDBException {
		if (columnFamilys.containsKey(columnFamilyName)) {
			return columnFamilys.get(columnFamilyName);
		} else {
			ColumnFamilyHandle handle = db.createColumnFamily(new ColumnFamilyDescriptor(columnFamilyName.getBytes(), COLUMN_FAMILY_OPTIONS));
			columnFamilys.put(columnFamilyName, handle);
			return handle;
		}
	}

	@Override
	public boolean containsKey(String keyspaces, String key) throws RocksDBException {
		byte[] bs = db.get(getColumnFamily(keyspaces), READ_OPTIONS, key.getBytes());
		if(bs != null) { return true; }
		return false;
	}

	@Override
	public byte[] get(String keyspaces, String key) throws RocksDBException {
		return db.get(getColumnFamily(keyspaces), READ_OPTIONS, key.getBytes());
	}

	@SuppressWarnings("resource")
	@Override
	public void forEach(String keyspaces, BiConsumer<String, byte[]> consumer) throws RocksDBException {
		Snapshot snapshot = db.getSnapshot();
		try {
			try(ReadOptions readOptions = new ReadOptions().setFillCache(true).setPinData(true).setTotalOrderSeek(true).setSnapshot(snapshot)){
				try(RocksIterator iterator = db.newIterator(getColumnFamily(keyspaces), readOptions)){
					for(iterator.seekToFirst(); iterator.isValid(); iterator.next()) {
						consumer.accept(new String(iterator.key()), iterator.value());
					}
				}
			}
		}finally {
			db.releaseSnapshot(snapshot);
			snapshot.close();
		}
	}
	
	@Override
	public void put(String keyspaces, String key, byte[] value) throws RocksDBException {
		db.put(getColumnFamily(keyspaces), WRITE_OPTIONS, key.getBytes(), value);
	}
	
	@Override
	public void putAll(String keyspaces, Map<String, byte[]> datas) throws RocksDBException {
		ColumnFamilyHandle handle = getColumnFamily(keyspaces);
		try(WriteBatch batch = new WriteBatch()){
			for (Entry<String, byte[]> entry : datas.entrySet()) {
				batch.put(handle, entry.getKey().getBytes(), entry.getValue());
			}
			db.write(WRITE_OPTIONS, batch);
		}
	}

	@Override
	public void remove(String keyspaces, String key) throws RocksDBException {
		db.delete(getColumnFamily(keyspaces), WRITE_OPTIONS, key.getBytes());
	}
	
	@Override
	public Set<String> keyspaces() throws RocksDBException {
		return columnFamilys.keySet();
	}
	@SuppressWarnings("resource")
	@Override
	public void clear(String keyspaces) throws RocksDBException {
		ColumnFamilyHandle handle = getColumnFamily(keyspaces);
		try(WriteOptions write = new WriteOptions(WRITE_OPTIONS).setSync(true)){
			byte[] start = new byte[]{0};
			byte[] end = new byte[]{127};
			db.deleteRange(handle, write, start, end);
			try(RocksIterator iterator = db.newIterator(handle)){
				try(WriteBatch batch = new WriteBatch()){
					for(iterator.seekToFirst(); iterator.isValid(); iterator.next()) {
						batch.delete(handle, iterator.key());
						if(batch.count() > 10000){
							db.write(write, batch);
							batch.clear();
						}
					}
					db.write(write, batch);
				}
			}
		}
	}
	
	public void clearAll() throws RocksDBException{
		for (Entry<String, ColumnFamilyHandle> entry : columnFamilys.entrySet()) {
			if(!entry.getKey().equals("default")){
				db.dropColumnFamily(entry.getValue());
			}
		}
		columnFamilys.clear();
	}
	
	@SuppressWarnings("resource")
	private void flush() {
		try {
			db.compactRange();
			db.flush(new FlushOptions().setWaitForFlush(true));
		} catch (RocksDBException e) {
			e.printStackTrace();
		}
	}
	
	@Override
	public void close() {
		columnFamilys.forEach((keyspaces, cf)->{ cf.close(); });
		columnFamilys.clear();
		flush();
		db.close();
		READ_OPTIONS.close();
		WRITE_OPTIONS.close();
		COLUMN_FAMILY_OPTIONS.close();
		DB_OPTIONS.close();
		OPTIONS.close();
	}
}
